

import os
import numpy as np
from PIL import Image
import math
from tqdm import tqdm  # 导入进度条库


def rotate_columns_independently(input_path, output_path, angle=10):
    # 读取图片并转换为numpy数组
    img = Image.open(input_path)
    img_array = np.array(img)

    # 获取图像尺寸
    if len(img_array.shape) == 3:  # 彩色图像
        height, width, channels = img_array.shape
        is_color = True
    else:  # 灰度图像
        height, width = img_array.shape
        is_color = False
        channels = 1
        img_array = img_array[:, :, np.newaxis]  # 增加通道维度

    # 创建空白的新图像（扩大尺寸以容纳旋转后的像素）
    scale_factor = 1.5
    new_height = int(height * scale_factor)
    new_width = int(width * scale_factor)

    # 计算中心偏移量
    offset_y = (new_height - height) // 2
    offset_x = (new_width - width) // 2

    # 创建新图像（白色背景）
    if is_color:
        new_img = np.zeros((new_height, new_width, channels), dtype=np.uint8) + 255
    else:
        new_img = np.zeros((new_height, new_width), dtype=np.uint8) + 255

    # 计算旋转角度（转换为弧度）
    theta = math.radians(angle)
    cos_theta = math.cos(theta)
    sin_theta = math.sin(theta)

    # 原图像中心点
    center_y = height // 2

    # 使用反向映射和双线性插值，添加进度条
    print(f"旋转处理中: {new_height}行 × {new_width}列...")
    for new_y in tqdm(range(new_height), desc="处理进度", unit="行"):
        for new_x in range(new_width):
            # 转换到原图像坐标系
            x0 = new_x - offset_x
            y0 = new_y - offset_y

            # 计算逆变换坐标
            # 公式推导：
            # 正向变换: new_x = col - (row - center_y)*sinθ
            #          new_y = center_y + (row - center_y)*cosθ
            # 逆变换解:
            if abs(cos_theta) > 1e-6:
                rel_y = (y0 - center_y) / cos_theta
            else:
                rel_y = 0  # 避免除以零

            col = x0 + sin_theta * rel_y
            row = center_y + rel_y

            # 检查是否在原图像范围内
            if 0 <= col < width and 0 <= row < height:
                # 双线性插值
                col0 = int(math.floor(col))
                col1 = min(col0 + 1, width - 1)
                row0 = int(math.floor(row))
                row1 = min(row0 + 1, height - 1)

                # 计算权重
                a = col - col0
                b = row - row0

                # 获取四个相邻像素
                if is_color:
                    p00 = img_array[row0, col0]
                    p01 = img_array[row0, col1]
                    p10 = img_array[row1, col0]
                    p11 = img_array[row1, col1]

                    # 双线性插值计算
                    interpolated = (
                        (1 - a) * (1 - b) * p00
                        + a * (1 - b) * p01
                        + (1 - a) * b * p10
                        + a * b * p11
                    )
                    new_img[new_y, new_x] = np.clip(interpolated, 0, 255).astype(
                        np.uint8
                    )
                else:
                    p00 = img_array[row0, col0, 0]
                    p01 = img_array[row0, col1, 0]
                    p10 = img_array[row1, col0, 0]
                    p11 = img_array[row1, col1, 0]

                    interpolated = (
                        (1 - a) * (1 - b) * p00
                        + a * (1 - b) * p01
                        + (1 - a) * b * p10
                        + a * b * p11
                    )
                    new_img[new_y, new_x] = np.clip(interpolated, 0, 255).astype(
                        np.uint8
                    )

    # 创建并保存新图像
    if not is_color:
        new_img = new_img[:, :, 0] if new_img.ndim == 3 else new_img  # 移除多余的维度
    result_img = Image.fromarray(new_img)
    result_img.save(output_path)
    print(f"旋转完成！结果已保存至: {output_path}")


# 使用示例
if __name__ == "__main__":
    # input_image = "1-2b.bmp"
    # output_image = "1-2b_rotated_columns.jpg"
    # rotate_columns_independently(input_image, output_image, angle=-30)

    # input_image = "1-3b.bmp"
    # output_image = "1-3b_rotated_columns.jpg"
    # rotate_columns_independently(input_image, output_image, angle=-30)
    
    file_dir = R"G:\idrs-project\bonan\bonan_core\白光图像"
    for i in range(1,20):
        input_image = os.path.join(file_dir, "1-{}b.bmp".format(i))
        output_image = os.path.join(file_dir, "1-{}b_rotated_columns.jpg".format(i))
        if(not os.path.exists(input_image)):
            continue
        rotate_columns_independently(input_image, output_image, angle=-30)
